package com.bestv.search.engine.compass.extend;

import com.bestv.search.common.util.StringConst;
import org.apache.lucene.document.Document;
import org.apache.lucene.search.HitCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.util.PriorityQueue;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author Michael Wang
 */

// $Id$
public class TypeTopDocCollector extends HitCollector {

    /**
     * The total number of hits the collector encountered.
     */
    protected int totalHits;
    protected int numHits;
    /**
     * The priority queue which holds the top-scoring documents.
     */
    protected PriorityQueue hq;
    /**
     * The priority queue which holds the top-scoring and charging documents.
     */
    protected PriorityQueue chargeHq;
    private ScoreDoc chargeReusableSD;
    private ScoreDoc freeReusableSD;
    private IndexSearcher indexSearcher;
    private Map<String, Integer> typeCount = new HashMap<String, Integer>();

    /**
     * Construct to collect a given number of hits.
     *
     * @param numHits       the maximum number of hits to collect
     * @param indexSearcher searcher instance to get document object
     */
    public TypeTopDocCollector(int numHits, IndexSearcher indexSearcher) {
        this(new TypeHitQueue(numHits));
        this.indexSearcher = indexSearcher;
        this.chargeHq = (new CategoryHitQueue(numHits));
        this.numHits = numHits;
    }

    /**
     * Constructor to collect the top-scoring documents by using the given PQ.
     *
     * @param hq the PQ to use by this instance.
     */
    protected TypeTopDocCollector(PriorityQueue hq) {
        this.hq = hq;
    }

    @Override
    public void collect(int doc, float score) {
        if (score > 0.0f) {
            totalHits++;
            int isCharge = 0;
            try {
                Document d = indexSearcher.doc(doc);
                String type = d.get("type");
                if (type == null || StringConst.EMPTY.equals(type) || StringConst.NULLSTRING.equalsIgnoreCase(type)) {
                    type = StringConst.OTHER_CHINESE;
                }

                isCharge = Integer.valueOf(d.get("isCharge"));
                if (typeCount.containsKey(type)) {
                    typeCount.put(type, typeCount.get(type) + 1);
                } else {
                    typeCount.put(type, 1);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }

            // 判断节目是否收费， 如果收费，放入收费列表中
            if (isCharge == 1) {
                if (chargeReusableSD == null) {
                    chargeReusableSD = new ScoreDoc(doc, score);
                } else if (score >= chargeReusableSD.score) {
                    // reusableSD holds the last "rejected" entry, so, if
                    // this new score is not better than that, there's no
                    // need to try inserting it
                    chargeReusableSD.doc = doc;
                    chargeReusableSD.score = score;
                } else {
                    return;
                }
                chargeReusableSD = (ScoreDoc) chargeHq.insertWithOverflow(chargeReusableSD);
            } else {
                if (freeReusableSD == null) {
                    freeReusableSD = new ScoreDoc(doc, score);
                } else if (score >= freeReusableSD.score) {
                    // reusableSD holds the last "rejected" entry, so, if
                    // this new score is not better than that, there's no
                    // need to try inserting it
                    freeReusableSD.doc = doc;
                    freeReusableSD.score = score;
                } else {
                    return;
                }
                freeReusableSD = (ScoreDoc) hq.insertWithOverflow(freeReusableSD);
            }
        }
    }

    /**
     * The total number of documents that matched this query.
     *
     * @return total hit count
     */
    public int getTotalHits() {
        return totalHits;
    }

    /**
     * The top-scoring hits.
     *
     * @return top scored documents
     */
    public TopDocs topDocs() {
        float maxScore = 0;
        // 合并 hq和chargeHq的doc到list中，用于排序
        List<CustomTopDoc> list = new ArrayList<CustomTopDoc>();
        int chargeHqSize = chargeHq.size();
        for (int i = chargeHqSize - 1; i >= 0; i--) {
            ScoreDoc scoreDoc = (ScoreDoc) chargeHq.pop();
            list.add(0, new CustomTopDoc(scoreDoc.doc, scoreDoc.score, 1));
            if (scoreDoc.score > maxScore) {
                maxScore = scoreDoc.score;
            }
        }

        int hqSize = hq.size();
        for (int i = hqSize - 1; i >= 0; i--) {
            ScoreDoc scoreDoc = (ScoreDoc) hq.pop();
            list.add(chargeHqSize, new CustomTopDoc(scoreDoc.doc, scoreDoc.score, 0));
            if (scoreDoc.score > maxScore) {
                maxScore = scoreDoc.score;
            }
        }

        // 获取前scoreDocs.length个doc
        ScoreDoc[] scoreDocs = new ScoreDoc[Math.min(list.size(), numHits)];
        for (int i = 0; i < scoreDocs.length; i++) {
            CustomTopDoc temp = list.get(i);
            scoreDocs[i] = new ScoreDoc(temp.getDoc(), temp.getScore());
        }

        return new TopDocs(totalHits, scoreDocs, maxScore);
    }

    public Map<String, Integer> getTypeCount() {
        return typeCount;
    }
}
